//============================================================================
// Name        : tree-match.cpp
// Author      : Long
// Version     :
// Copyright   : Your copyright notice
// Description : Hello World in C++, Ansi-style
//============================================================================

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <armadillo>

#include "line_search.h"
#include "gradient_descent.h"
#include "accelerated_gradient.h"
#include "simple_set.h"
#include "projected_gradient_descent.h"
#include "accelerated_projected_gradient.h"
#include "admm.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

int n, d;
std::string gin, fin, out;

void
process_options(int argc, char** argv);

void
usage(int exitcode = 0)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

/** template for a function */
class HalfLengthSquareFunction
{
public:
  typedef double range_type;
  typedef arma::vec domain_type;
public:
  static const bool has_value_ = true;
  static const bool has_gradient_ = true;
  static const bool has_both_ = true;
public:
  range_type
  value(const domain_type& x)
  {
    return 0.5 * arma::dot(x, x);
  }

  domain_type
  gradient(const domain_type& x)
  {
    return x;
  }

  void
  value_gradient(const domain_type& x, range_type &value, domain_type& gradient)
  {
    value = this->value(x);
    gradient = this->gradient(x);
  }
};

class WeightedSquareFunction
{
public:
  typedef double range_type;
  typedef arma::vec domain_type;
public:
  static const bool has_value_ = true;
  static const bool has_gradient_ = true;
  static const bool has_both_ = true;
protected:
  domain_type weight_;
public:
  WeightedSquareFunction(const domain_type& weight) :
    weight_(weight)
  {
  }
  range_type
  value(const domain_type& x)
  {
    return arma::dot(weight_, x % x);
  }

  domain_type
  gradient(const domain_type& x)
  {
    return 2 * (weight_ % x);
  }

  void
  value_gradient(const domain_type& x, range_type &value, domain_type& gradient)
  {
    value = this->value(x);
    gradient = this->gradient(x);
  }
};

class NegativeEntropyFunction
{
public:
  typedef double range_type;
  typedef arma::vec domain_type;
public:
  static const bool has_value_ = true;
  static const bool has_gradient_ = true;
  static const bool has_both_ = true;
protected:
public:
  NegativeEntropyFunction()
  {
  }

  range_type
  value(const domain_type& x)
  {
    double s = 0.0;
    for (unsigned int i = 0; i < x.n_elem; i++)
      {
        if (x[i] > 1e-10)
          s += x[i] * log(x[i]);
        else if (x[i] < 0) return std::numeric_limits<double>::infinity();
      }
    return s;
  }

  domain_type
  gradient(const domain_type& x)
  {
    return arma::log(x) + 1;
  }

  void
  value_gradient(const domain_type& x, range_type &value, domain_type& gradient)
  {
    value = this->value(x);
    gradient = this->gradient(x);
  }
};

class LinearFunction
{
public:
  typedef double range_type;
  typedef arma::vec domain_type;
public:
  static const bool has_value_ = true;
  static const bool has_gradient_ = true;
  static const bool has_both_ = true;
protected:
  domain_type weight_;
public:
  LinearFunction()
  {
  }

  void
  setWeight(const domain_type& weight)
  {
    weight_ = weight;
  }

  range_type
  value(const domain_type& x)
  {
    return arma::dot(weight_, x);
  }

  domain_type
  gradient(const domain_type& x)
  {
    return weight_;
  }

  void
  value_gradient(const domain_type& x, range_type &value, domain_type& gradient)
  {
    value = this->value(x);
    gradient = this->gradient(x);
  }
};

void FirstOrderMethods()
{
  //typedef HalfLengthSquareFunction func_type;
  //typedef WeightedSquareFunction func_type;
  //typedef NegativeEntropyFunction func_type;
  typedef LinearFunction func_type;
  typedef func_type::domain_type domain_type;
  typedef func_type::range_type range_type;
  typedef LineSearch<func_type> optim_type;

  //  domain_type w(2);
  //  w[0] = 1;
  //  w[1] = 1;
  //  func_type f(w);
  func_type f;
  domain_type x0 = arma::randn<domain_type>(2) * 10, x, w = arma::randu<domain_type>(2)*10-5;
  w[0] = -1; w[1]=-1.0001;
  f.setWeight(w);

  std::cout << "w=\n" << w << std::endl;

  //  range_type f0 = f.value(x0);
  //
  //  std::cout << "x0=\n" << x0 << " f0=" << f0 << std::endl;

  //  double step = 1.0;
  //  domain_type grad = f.gradient(x0), dir = -grad;
  //  optim_type op(f, x0, x, f0, grad, dir, step);
  //  bool ok = op.search();
  //
  //  std::cout << "line search " << (ok ? "succeeded" : "failed") << "\n" << x0
  //      << "-->\n" << x << "\nf0=" << f0 << " fx=" << op.fx() << std::endl;

  //  typedef GradientDescent<func_type> gd_type;
  //  gd_type gd(f, x0, x);
  //  ok = gd.search();
  //  std::cout << "gradient descent " << (ok ? "succeeded" : "failed") << "\n" << x0
  //      << "-->\n" << x << "\nf0=" << f0 << " fx=" << gd.fx() << std::endl;

  //  typedef AcceleratedGradient<func_type> agd_type;
  //  agd_type agd(f, x0, x);
  //  ok = agd.search();
  //  std::cout << "accelerated gradient " << (ok ? "succeeded" : "failed") << "\n" << x0
  //      << "-->\n" << x << "\nf0=" << f0 << " fx=" << agd.fx() << std::endl;

  //x0.zeros();
  //  x0[0] = 0.1; x0[1] = 0.9;
  //  x = cs.projection(x0);
  //  std::cout << "project to canonical simplex " << "\n" << x0
  //      << "-->\n" << x << std::endl;

  //typedef CanonicalSimplex set_type;
  typedef StandardSimplex set_type;

  set_type cs;
  x0.randu();
  x0 = cs.projection(x0-0.5);

  //x0 /= sum(x0);

  typedef ProjectedGradientDescent<func_type, set_type> pgd_type;
  pgd_type pgd(f, cs, x0, x);
  bool ok = pgd.search();
  std::cout << "projected gradient descent " << (ok ? "succeeded" : "failed")
      << "\n" << x0 << "-->\n" << x << "\nf0=" << f.value(x0) << " fx="
      << pgd.fx() << std::endl;

  typedef AcceleratedProjectedGradient<func_type, set_type> apg_type;
  apg_type apg(f, cs, x0, x);
  ok = apg.search();
  std::cout << "accelerated projected gradient descent " << (ok ? "succeeded"
      : "failed") << "\n" << x0 << "-->\n" << x << "\nf0=" << f.value(x0)
      << " fx=" << apg.fx() << std::endl;
}

void ADMM()
{
  arma::mat A(2,2);
  arma::vec b(2);

  A.at(0,0) = 2;
  A.at(0,1) = 4;
  A.at(1,0) = 3;
  A.at(1,1) = 6;
  b.at(0) = 6;
  b.at(1) = 9;

  Lasso lasso(A, b, 5.0);
  lasso.search();

  std::cout << "lasso by admm: x=\n" << lasso.x() << "fx=" << lasso.fx() << std::endl;
}

int
main(int argc, char** argv)
{
  process_options(argc, argv);

  if (vm.count("help")) usage();

  srand(time(NULL));

  //FirstOrderMethods();
  ADMM();

  return 0;
}

void
process_options(int argc, char** argv)
{
  desc.add_options()("help", "produce this usage")
  //    ("gin", po::value<std::string>(&gin)->default_value(""), "input file (factor graph)")
  //    ("fin", po::value<std::string>(&fin)->default_value(""), "input file (factor functions)")
  //    ("out", po::value<std::string>(&out)->default_value(""), "output file")
  ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}
